import pandas as pd
import numpy as np
import random
from datetime import datetime, timedelta


class DataGenerator:
    """数据生成器，负责生成所有模拟数据"""

    def __init__(self):
        # 基础数据配置
        self.regions = ["华东区", "华北区", "华南区", "西北区", "西南区"]
        self.managers = ["张三", "李四", "王五", "赵六", "钱七"]
        self.business_types = ["再保财险", "再保人身险", "人身险", "财险"]
        self.contract_statuses = ["有效", "已到期", "已退保", "犹豫期退保"]

    def generate_regions_data(self):
        """生成区域和经理数据"""
        data = {
            "region_id": range(1, len(self.regions) + 1),
            "region_name": self.regions,
            "manager_name": self.managers,
            "manager_phone": [f"138{random.randint(10000000, 99999999)}" for _ in range(len(self.regions))]
        }
        return pd.DataFrame(data)

    def generate_contracts_data(self, num_records=1000):
        """生成合同数据"""
        contract_ids = [f"C{20230000 + i}" for i in range(1, num_records + 1)]

        # 随机生成日期（过去2年）
        start_date = datetime(2022, 1, 1)
        end_date = datetime(2023, 12, 31)
        date_range = [start_date + timedelta(days=random.randint(0, (end_date - start_date).days))
                      for _ in range(num_records)]

        # 随机生成保险期限（年）
        terms = [random.choice([1, 2, 3, 5, 10, 15, 20, 30]) for _ in range(num_records)]

        # 生成数据
        data = {
            "contract_id": contract_ids,
            "region_id": [random.randint(1, len(self.regions)) for _ in range(num_records)],
            "business_type": [random.choice(self.business_types) for _ in range(num_records)],
            "is_lead_reinsurer": [random.choice([True, False]) for _ in range(num_records)],
            "insurance_type": [random.choice(["长期险", "短期险"]) for _ in range(num_records)],
            "customer_type": [random.choice(["团险", "个险"]) for _ in range(num_records)],
            "policy_type": [random.choice(["首年期缴", "趸缴"]) for _ in range(num_records)],
            "term": terms,  # 保险期限（年）
            "issue_date": date_range,
            "status": [random.choice(self.contract_statuses) for _ in range(num_records)],
            "premium_income": [round(random.uniform(1000, 1000000), 2) for _ in range(num_records)],
            "reinsurance_premium": [round(random.uniform(500, 500000), 2) for _ in range(num_records)],
            "estimated_premium": [round(p * random.uniform(0.9, 1.1), 2)
                                  for p in [round(random.uniform(500, 500000), 2) for _ in range(num_records)]],
            "new_business_value": [round(random.uniform(100, 100000), 2) for _ in range(num_records)],
            "effective_business_value": [round(random.uniform(500, 500000), 2) for _ in range(num_records)],
            "free_look_withdrawal": [round(random.uniform(0, 50000), 2) if random.random() < 0.1 else 0
                                     for _ in range(num_records)]
        }

        # 转换为DataFrame并添加年份和季度
        df = pd.DataFrame(data)
        df["year"] = df["issue_date"].dt.year
        df["quarter"] = df["issue_date"].dt.quarter
        df["period"] = df["year"].astype(str) + "Q" + df["quarter"].astype(str)

        return df

    def generate_assets_data(self):
        """生成资产数据"""
        periods = []
        region_ids = []
        start_assets = []
        end_assets = []

        # 为每个区域生成8个季度的资产数据（2022Q1到2023Q4）
        for region_id in range(1, len(self.regions) + 1):
            # 初始资产
            base_asset = random.uniform(10000000, 50000000)

            for year in [2022, 2023]:
                for quarter in [1, 2, 3, 4]:
                    periods.append(f"{year}Q{quarter}")
                    region_ids.append(region_id)

                    # 资产有一定增长
                    start_asset = base_asset * (1 + random.uniform(-0.02, 0.05))
                    end_asset = start_asset * (1 + random.uniform(0.01, 0.1))

                    start_assets.append(round(start_asset, 2))
                    end_assets.append(round(end_asset, 2))

                    base_asset = end_asset

        data = {
            "period": periods,
            "region_id": region_ids,
            "start_asset": start_assets,
            "end_asset": end_assets
        }

        return pd.DataFrame(data)

    def generate_time_periods(self):
        """生成时间周期列表"""
        periods = []
        for year in [2022, 2023]:
            for quarter in [1, 2, 3, 4]:
                periods.append(f"{year}Q{quarter}")
        return periods
